package com.example.partition;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.springframework.batch.core.partition.support.Partitioner;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.HashMap;
import java.util.Map;

/**
 * 定义数据的分片规则
 *
 * @author lixing
 */
@Slf4j
public class RangePartitioner implements Partitioner {
    public static final String MINRECODE = "_minRecode";
    public static final String MAXRECODE = "_maxRecode";
    public static final String MIN_SELECT_PATTERN = "select min(id) as id from user";
    public static final String MAX_SELECT_PATTERN = "select max(id) as id from user";

    @Autowired
    private JdbcTemplate jdbcTemplate;

    @Override
    public Map<String, ExecutionContext> partition(int gridSize) {
        log.info("partition called gridsize= " + gridSize);

        Map<String, ExecutionContext> result = new HashMap<String, ExecutionContext>();

        //取到id最小值
        int min = (int) jdbcTemplate.queryForMap(MIN_SELECT_PATTERN).get("id");
        //取到id最大值
        int max = (int) jdbcTemplate.queryForMap(MAX_SELECT_PATTERN).get("id");

        //计算每一片的数据条数
        int targetSize = (max - min) / gridSize + 1;
        int number = 0;
        int start = min;
        int end = start + targetSize - 1;;

        //构建每一片的ExecutionContext，里面设置的值可以在@StepScope 取到
        while (start <= max) {
            ExecutionContext context = new ExecutionContext();
            if (end >= max) {
                end = max;
            }
            number++;
            context.putInt(MINRECODE, start);
            context.putInt(MAXRECODE, end);
            context.putString("name", "Thread" + number);

            start += targetSize;
            end += targetSize;
            result.put("partition" + number, context);
        }
        System.out.println(JSON.toJSONString(result));
        return result;

        /*for (int i = 1; i <= gridSize; i++) {
            ExecutionContext value = new ExecutionContext();

            System.out.println("\nStarting : Thread" + i);
            System.out.println("fromId : " + fromId);
            System.out.println("toId : " + toId);

            value.putInt("fromId", fromId);
            value.putInt("toId", toId);

            // give each thread a name, thread 1,2,3
            value.putString("name", "Thread" + i);

            result.put("partition" + i, value);

            fromId = toId + 1;
            toId += range;

        }
        return result;*/
    }
}